import re
from langchain.prompts import SystemMessagePromptTemplate, ChatPromptTemplate
from langchain_experimental.plan_and_execute import (
    PlanAndExecute,
    load_agent_executor,
    load_chat_planner,
)
from langchain_openai import ChatOpenAI
import os
from langchain.agents import tool
import json

os.environ["LANGSMITH_TRACING"] = "true"
# os.environ["LANGSMITH_API_KEY"] = getpass.getpass()
os.environ["LANGSMITH_API_KEY"] = "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"

llm = ChatOpenAI(
    # openai_api_key="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",  # gtihub
    openai_api_key="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",  # 谷歌
    openai_api_base="https://openrouter.ai/api/v1",
    # model_name="deepseek/deepseek-chat-v3-0324:free",
    model_name="qwen/qwq-32b:free",
)


plan_tmp = """\
    首先请你理解这个问题，然后设计解决这个问题的计划，设计计划过程中必须遵守以下规则：
    1、计划的表头必须使用“plan:”开始，然后是步骤的编号列表
    2、以最少的计划步骤准确的完成计划设计
    3、设计计划时优先考虑使用工具列表中的工具完成计划任务，工具列表：{Tools}
    4、设计计划时需考虑计划的依赖性和连贯性，尽量使用原始问题中的相关信息调用对应工具完成计划任务
    5、调用工具时严格遵守工具的参数要求，具体信息参考工具列表详细描述
    6、计划步骤中严禁出现重复计划或者相似的重复计划
    7、计划的最后一步需考虑之前的计划是否能解决原始问题
    8、在计划的末尾必须使用“<END_OF_PLAN>”表示计划结束
    9、相关工具返回结果中“<END_RUN>”表示结束当前全部计划
    
    以下为提供设计计划的工具列表详细描述：
    {Tools_Describe}
"""


@tool()
def overview(*args, **kwargs):
    """
    js概览组件
    描述：当需要展示关键指标摘要时使用
    """

    res = {
        "title": "概览的标题",
        "type": "概览",
        "data": [{"label": "标签", "value": "数值"}],
    }

    return json.dumps(res, ensure_ascii=False)


@tool()
def barChart(*args, **kwargs) -> str:
    """
    js柱状图组件 - 用于比较不同类别间的数值差异

    参数:
        title: str - 图表标题
        categories: List[str] - 类别列表
        values: List[float] - 对应数值列表

    返回:
        JSON字符串格式的柱状图数据
    """
    data = {
        "title": "网络安全事件统计",
        "type": "barChart",
        "data": [
            {"label": "高危事件", "value": 5},
            {"label": "中危事件", "value": 6},
            {"label": "低危事件", "value": 2},
        ],
    }

    return json.dumps(data, ensure_ascii=False)
    # return "柱状图描述"


@tool(description="组件数据拼接工具，用于将多个可视化组件的数据合并为统一的JSON格式")
def buidlJSon(*args, **kwargs) -> str:
    """
    将多个可视化组件的数据智能组合成标准化的仪表板数据结构。

    返回数据结构规范：
    {
        "header": {
            "title": "自动提炼最核心的主题作为标题(要求：不超过10个汉字)",
            "description": "自动生成1句简洁的摘要(要求：20-30字，概括数据核心价值)"
        },
        // 各个组件
        "content": [
            {
                "title": "图表标题(自动提取原数据特征)",
                "type": "chartType",  // 支持 barChart/lineChart/pieChart 等
                "data": [  // 标准化数据格式
                    {"label": "分类名称", "value": 数值},
                    ...
                ]
            },
            ...
        ]
    }

    智能处理要求：
    1. 自动识别输入数据的核心特征生成有意义的标题
    2. 当多个图表存在关联性时，应生成反映整体洞察的描述
    3. 保持数据一致性，确保数值单位和分类标准统一
    4. 对异常数据结构提供容错处理
    """


# 工具列表
tools = [overview, barChart, buidlJSon]
# tools = [test]
tool_strings = []
for t in tools:
    args_schema = re.sub("}", "}}", re.sub("{", "{{", str(t.args)))
    tool_strings.append(f"{t.name}: {t.description}, args: {args_schema}")
    formatted_tools = "\n".join(tool_strings)
    tool_names = ", ".join([t.name for t in tools])

sys_prompt_tmp = SystemMessagePromptTemplate.from_template(plan_tmp)
chat_sys_tmp = ChatPromptTemplate.from_messages([sys_prompt_tmp])
sys_prompt = chat_sys_tmp.format(Tools=tool_names, Tools_Describe=formatted_tools)

planner = load_chat_planner(llm, system_prompt=sys_prompt)
executor = load_agent_executor(llm, tools, verbose=True)


summary = """
{
    "网络安全摘要信息": {
        "资产网络攻击统计相关数据": {
            "资产ID": "0fd4b9d96d6f5ecab6c6a8b6809f5ae3",
            "资产名称": "taomingwei-node3-ubuntu",
            "总事件数量": 13,
            "告警时间范围": {
                "第一条告警时间": "2022年4月3日09:52:39",
                "最新告警时间": "2025年4月8日09:12:25"
            }
        }
}
"""


try:
    agent = PlanAndExecute(planner=planner, executor=executor, verbose=True)
    input = """\
        请根据提供的输入内容生成一个网站的网页JSON

        输入内容：
        {summary}

    """
    input = input.format(summary=summary)
    res = agent.invoke(input, verbose=True)

    print("res", res)
    # print("输出结果", res.content)
except Exception as e:
    print("异常", e, input)
